import random, sys, time
from pathlib import Path
from typing import Sequence
from sumolib.net import readNet, Net
from feasytools import ArgChecker
from fpowerkit import Grid

HELP_STR = f"""
充电站生成程序
{sys.argv[0]} -d <SUMO配置文件夹> [--slots-cs <快充站充电桩数量>] [--slots-pk <路边慢充充电桩数量>] [--psell <售电价格>] \\
[--n-cs <快充站数量>] [--n-pk <路边慢充站数量>] [--n-bus <母线数量>] [--buses <使用的母线>] [--seed <随机化种子>] [--non-cs] [--non-pk] [-h/--help]
  h/help: 显示此帮助信息
  d: SUMO配置文件夹
  slots-cs: 快充站充电桩数量, 默认20
  slots-pk: 路边慢充充电桩数量, 默认20
  psell: 售电价格, 默认1元/kWh
  n-cs: 使用的快充站数量, 默认使用所有
  n-pk: 使用的路边慢充站数量, 默认使用所有
  n-bus: 使用的母线数量, 默认使用所有, 该选项不能与buses一起使用
  buses: 使用的母线, 默认使用所有, 例如"B1,B2,B3", 该选项不能与n-bus一起使用
  seed: 随机化种子, 不指定则使用当前时间(纳秒)
  non-cs: 不生成快充站文件
  non-pk: 不生成慢充站文件
"""


def __checkFile(file: str):
    p = Path(file)
    if p.exists():
        i = 1
        while True:
            p = Path(file + f".bak{i}")
            i += 1
            if not p.exists():
                break
        Path(file).rename(str(p))


def __detect(file: str):
    rnet: Net = readNet(file)
    edges: list[str] = [e.getID() for e in rnet.getEdges()]
    cs: list[str] = [
        e for e in edges if e.upper().startswith("CS") and not e.lower().endswith("rev")
    ]
    pk: list[str] = [e for e in edges if not e.upper().startswith("CS")]
    return cs, pk


def __generate(
    file: str, cs_names: list[str], buses: list[str], psell: float, slots: int
):
    __checkFile(file)
    fp = open(file, "w")
    fp.write("<root>")
    for cs in cs_names:
        fp.write(f'<cs name="{cs}" slots="{slots}" bus="{random.choice(buses)}">\n')
        fp.write(f"  <pbuy>\n")
        fp.write(f'    <item btime="0" price="1.50" />\n')
        for d in range(0, 8):
            t = [0, 0, 0, 0, 0]
            p = [0.0, 0.0, 0.0, 0.0, 0]
            t[0] = random.choice([0, 1]) if d > 0 else 1
            p[0] = random.uniform(1.0, 1.1)
            t[1] = random.choice([6, 7, 8])
            p[1] = random.uniform(1.8, 2.1)
            t[2] = random.choice([10, 11])
            p[2] = random.uniform(1.4, 1.6)
            t[3] = random.choice([15, 16, 17])
            p[3] = random.uniform(1.7, 2.0)
            t[4] = random.choice([19, 20])
            p[4] = random.uniform(1.5, 1.7)
            for tx, px in zip(t, p):
                fp.write(f'    <item btime="{d*86400+tx*3600}" price="{px:.3}" />\n')
        fp.write(f"  </pbuy>\n")
        fp.write(f"  <psell>\n")
        fp.write(f'    <item btime="0" price="{psell}" />\n')
        fp.write(f"  </psell>\n</cs>\n")
    fp.write("</root>")
    fp.close()


def generate_cspk(
    seed: int,
    root: str,
    n_bus: int,
    new_buses: Sequence[str],
    psell: float,
    slots_cs: int,
    slots_pk: int,
    n_cs: int = -1,
    n_pk: int = -1,
    non_cs: bool = False,
    non_pk: bool = False,
    silent: bool = False,
):
    def elprint(con: str):
        if not silent:
            print(con)

    random.seed(seed)
    root = root.rstrip("/\\")
    pos = root.rfind("/")
    if pos >= 0:
        root = root + "/" + root[pos + 1 :]
    else:
        root = root + "/" + root
    cs, pk = __detect(root + ".net.xml")
    buses = [x.ID for x in Grid.fromFile(root + ".grid.zip").Buses]
    if "B0" in buses:
        buses.remove("B0")
    elprint(f"检测到可用母线{len(buses)}个: {buses}")
    if n_bus > 0:
        buses = random.sample(buses, n_bus)
        elprint(f"  选中{len(buses)}个母线: {buses}")
        if len(new_buses) > 0:
            elprint("警告: 指定的母线选项--buses将被忽略")
    else:
        if len(new_buses) > 0:
            new_buses = [x.strip() for x in new_buses]
            if not set(new_buses).issubset(set(buses)):
                raise ValueError(f"错误: 指定的母线集{new_buses}不是可用母线的子集")
            buses = new_buses

    flag = False
    if not non_cs:
        elprint(f"检测到可用快充站{len(cs)}个: {cs}")
        if n_cs > 0:
            cs = random.sample(cs, n_cs)
            elprint(f"  选中{len(cs)}个快充站: {cs}")
        __generate(root + ".cs.xml", cs, buses, psell, slots_cs)
        elprint("  已生成快充站文件")
        flag = True

    if not non_pk:
        elprint(f"检测到可用路边慢充站{len(pk)}个: {pk}")

        if n_pk > 0:
            pk = random.sample(pk, n_pk)
            elprint(f"  选中{len(pk)}个快充站: {pk}")
        __generate(root + ".pk.xml", pk, buses, psell, slots_pk)
        elprint("  已生成路边慢充文件")
        flag = True

    if not flag:
        elprint("警告: 未生成任何数据")
    else:
        elprint("操作完成")


if __name__ == "__main__":
    params = ArgChecker()
    if params.pop_bool("h"):
        print(HELP_STR)
        sys.exit()

    seed = params.pop_int("seed", time.time_ns())
    root = params.pop_str("d")  # 12nodes
    n_bus = params.pop_int("n-bus", 0)
    new_buses = params.pop_str("buses", "").split(",")
    if len(new_buses) == 1 and len(new_buses[0]) == 0:
        new_buses = []
    psell = params.pop_float("psell", 1.0)
    slots_cs = params.pop_int("slots-cs", 20)
    slots_pk = params.pop_int("slots-pk", 20)
    n_cs = params.pop_int("n-cs", 0)
    n_pk = params.pop_int("n-pk", 0)
    non_cs = params.pop_bool("non-cs")
    non_pk = params.pop_bool("non-pk")
    generate_cspk(
        seed,
        root,
        n_bus,
        new_buses,
        psell,
        slots_cs,
        slots_pk,
        n_cs,
        n_pk,
        non_cs,
        non_pk,
    )
